import paramiko
import threading
from loguru import logger
import traceback
import re
import time
import os
import select


from global_config.global_env import GLO_VAR


global_commands = 0

class SSHMulti(threading.Thread):
    '''
    继承Thread实现paramiko多线程并发
    '''
    def __init__(self, host_ip, username, password, command, port=22, **kwargs):
        '''
        :param host_ip: 主机IP
        :param username: 用户名
        :param password: 密码
        :param command: 具体命令
        :param port: 端口号默认22
        :param kwargs:
        '''
        super().__init__()
        self.host_ip = host_ip
        self.username = username
        self.password = password
        self.command = command
        self.port = port
        self.result, self.error = None, None
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(hostname=self.host_ip, port=int(self.port), username=self.username, password=self.password)

    def run(self):
        '''
        ssh定义部分需放init中，run内处理命令发送逻辑
        '''
        global global_commands
        global_commands += 1
        # transport = paramiko.Transport(self.host_ip, self.port)
        # transport.banner_timeout = 300
        print('*' * 144)
        print('the %d SSH command:    %s' % (global_commands, self.command), end='\n\n')
        stdin, stdout, stderr = self.ssh.exec_command(self.command)
        self.error = [line for line in stderr]
        self.result = [line for line in stdout]
        print('the %d SSH command result:' % global_commands)
        [print(i, end='') for i in self.result]
        print()
        print('*' * 144)
        self.ssh.close()


def ssh_command(host_ip, username, password, command, port=22, ischeck=0, **kwargs):
    '''
    实现ssh会话连接，发送命令后返回结果
    :param host_ip: 主机IP
    :param username: 用户名
    :param password: 密码
    :param command: 具体命令
    :param port: 端口号默认22
    :param ischeck: 存在错误输出时是否处理，0--忽略，1--报异常
    :param kwargs:
    return：命令的结果，报错信息，列表
    '''
    global global_commands
    global_commands += 1
    # transport = paramiko.Transport(host_ip, port)
    # transport.banner_timeout = 300
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname=host_ip, port=int(port), username=username, password=password)
    except:
        traceback.print_exc()
        raise BaseException("ssh_command error")
    # print('*'*144)
    logger.info('>'*144)
    logger.info(f'the {global_commands} SSH command({host_ip}):    {command}')
    stdin, stdout, stderr = ssh.exec_command(command)
    err_list = [line for line in stderr]
    res_list = [line for line in stdout]
    logger.info(f'the {global_commands} SSH command result:')
    [logger.info(i.strip()) for i in res_list]
    [logger.info(i.strip()) for i in err_list]
    logger.info('<' * 144)
    ssh.close()
    if ischeck and err_list:
        logger.info(f"{command} 命令错误信息：")
        for i in err_list:
            logger.info(i)
        raise BaseException("ssh_connect下发命令 {} error，请确认".format(command))
    return res_list, err_list


def ssh_commands(host_ip, username, password, commands, port=22, ischeck=0, **kwargs):
    '''
    ssh会话连接，连续发送多条命令并返回结果
    :param host_ip: 主机IP
    :param username: 用户名
    :param password: 密码
    :param commands: 命令列表
    :param port: 端口号默认22
    :param ischeck: 存在错误输出时是否处理，0--忽略，1--报异常
    :param kwargs:
    return：命令的结果，报错信息，列表
    '''
    global global_commands

    # transport = paramiko.Transport(host_ip, port)
    # transport.banner_timeout = 300
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(hostname=host_ip, port=int(port), username=username, password=password)
    except:
        traceback.print_exc()
        raise BaseException("ssh_commands error")

    ret_list = []
    for command in commands:
        global_commands += 1
        logger.info('-'*144)
        logger.info(f'the {global_commands} SSH command({host_ip}):    {command}')
        stdin, stdout, stderr = ssh.exec_command(command)
        err_list = [line.strip() for line in stderr]
        res_list = [line.strip() for line in stdout]
        logger.info(f'the {global_commands} SSH command result:')
        [logger.info(i.strip()) for i in res_list]
        # print('*' * 144)
        logger.info('-' * 144)
        ret_list.append((res_list, err_list))
        if ischeck and err_list:
            logger.info(f"{command} 命令错误信息：")
            for i in err_list:
                logger.info(i)
            raise BaseException("ssh_connect下发命令 {} error，请确认".format(command))
    ssh.close()
    return ret_list



def ssh_interactive(host_ip, username, password, commands, re_compiles, port=22, timeout=10, **kwargs):
    """
    ssh连接实现发送交互式命令，返回结果（命令最长响应时间不可高于timeout）
    :param host_ip:远程主机ip
    :param username:用户名
    :param password:密码
    :param commands: 交互命令列表
    :param re_compiles:用于识别命令执行完毕的已编译re模式对象，需与commands对应
    :param port:端口
    :param timeout:每个命令等待的最大时间
    :param kwargs:
    :return: 命令执行结果，列表
    """
    # commands和conditions参数检查
    if not isinstance(commands, (list, tuple)) or not isinstance(re_compiles, (list, tuple)):
        raise BaseException(f"must be list or tuple!\ncommands:{commands}\nconditions:{re_compiles}")
    if len(commands) != len(re_compiles):
        raise BaseException(f"commands and conditions not compare!\ncommands:{commands}\nconditions:{re_compiles}")
    # 加入会话登录检查流程兼容远程主机响应速度低下问题
    commands = ["", ] + list(commands)
    re_compiles = [re.compile(r"[\s\S]*\]#", re.S), ] + list(re_compiles)
    # ------------------------
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    # 建立连接
    try:
        ssh.connect(hostname=host_ip, port=int(port), username=username, password=password)
        chan = ssh.invoke_shell()
        chan.set_combine_stderr(True)
        # print(dir(chan))
    except paramiko.SSHException:
        logger.info("SSH send command fail! SSHException")
        logger.info(traceback.format_exc())
        return False, traceback.format_exc()

    result = []
    for i, command in enumerate(commands):
        logger.info('-'*74)
        logger.info(f"command: {command}")
        chan.send(command)
        recv = ""
        start_time = time.time()
        while True:
            # print(recv)
            if time.time() - start_time >= timeout:
                error_info = "ssh_interactive timeout!"
                logger.info(error_info)
                logger.info("recv: ")
                logger.info(">"*74)
                logger.info(f"{recv}")
                logger.info("<" * 74)
                return False, error_info
                # raise BaseException("timeout!")
            r, _, _ = select.select([chan, ], [], [], 1)
            if not r:
                continue
            _recv = chan.recv(1024000)
            # err = chan.recv_stderr(1024000)
            # print(err)
            recv += _recv.decode("utf-8")
            print(f"match recv: {recv} \n#################################")
            if re_compiles[i].match(recv):
                logger.info(">" * 74)
                logger.info(f"output: \n{recv}")
                logger.info("<" * 74)
                result.append(recv)
                break
            # print(_recv)
            time.sleep(0.1)
    chan.close()
    ssh.close()
    return True, result



def ssh_interactive_yield(host_ip, username, password, port=22, timeout=30, linux=True, **kwargs):
    """
    ssh连接状态生成器，返回结果（命令最长响应时间不可高于timeout）,支持中文字符
    :param host_ip:远程主机ip
    :param username:用户名
    :param password:密码
    :param port:端口
    :param timeout:每个命令等待的最大时间
    :param linux: 远程主机为linux时，将识别linux终端控制字符，截取真实字串
    :param kwargs:
    :return: 命令执行结果，列表
    """
    linux_clear_char = b"\x1b[0K"  # linux标准清理字符，清除整个recv之前接收的字符

    def get_real_rsp_on_linux(raw_bytes: bytes):  # 识别linux清理字符并向后截取真实打印字符
        # return raw_bytes[raw_bytes.rfind(linux_clear_char)+4:]
        return raw_bytes.rsplit(linux_clear_char)[-1]

    # 强制重载paramiko
    # importlib.reload(paramiko)
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    # 建立连接
    try:
        ssh.connect(hostname=host_ip, port=int(port), username=username, password=password)
        chan = ssh.invoke_shell()
        chan.set_combine_stderr(True)
        # print(dir(chan))
    except paramiko.SSHException:
        logger.info("SSH send command fail! SSHException")
        logger.info(traceback.format_exc())
        return False, traceback.format_exc()

    result = True
    output = ""  # 转码utf后的实际字符
    while True:
        command, comp, timeout = yield result, output
        if command is None:  # 收到空命令认为会话中止，释放会话
            chan.close()
            ssh.close()
            break
        logger.info(f"command:  {command}re_ptn :  {comp}")
        chan.send(command)
        recv = b""
        start_time = time.time()
        while True:
            # print(recv)
            # 超时返回一个失败信息
            if time.time() - start_time >= timeout:
                logger.info(command)
                error_info = "ssh_interactive timeout!"
                logger.info("recv: ")
                logger.info(f"{recv}")
                logger.info(error_info)
                result, output = False, error_info
                break
            r, _, _ = select.select([chan, ], [], [], 1)
            if not r:
                continue
            _recv = chan.recv(1024000)
            # print("_recv:", _recv)
            recv += _recv
            if linux:
                recv = get_real_rsp_on_linux(recv)
            try:
                output = recv.decode()
            except:
                pass
            else:
                if comp.match(output):
                    print("recv:", recv)
                    result = True
                    # output = recv
                    # logger.info(f"ssh_interactive_yield output:\n{output}\n{'-'*74}")
                    break
            # print(_recv)
            time.sleep(0.5)




def sftp_upload_file(host_ip, username, password, remote_path, local_path, port=22,**kwargs):
    '''
    本地向远程主机上传文件
    :param host_ip: 主机IP
    :param username: 用户名
    :param password: 密码
    :param remote_path: 远程文件路径
    :param local_path: 待上传的本地文件路径
    :param port: sftp协议端口号
    :param kwargs:
    :return:
    '''
    # 创建SSH客户端
    t = paramiko.Transport((host_ip, port))
    t.connect(username=username, password=password)
    sftp = paramiko.SFTPClient.from_transport(t)
    logger.info(f"{local_path} {remote_path}")
    _mkdirs(sftp, os.path.dirname(remote_path))
    print(sftp.put(local_path, remote_path))
    t.close()


def sftp_down_file(host_ip, username, password, server_path, local_path, port=22):
    '''
    windows从linux客户机下载指定文件
    :param host_ip: 主机IP
    :param username: 用户名
    :param password: 密码
    :param server_path: linux客户端文件路径
    :param local_path: 下载到本地的文件路径
    :param kwargs:
    :return:
    '''
    t = paramiko.Transport((host_ip, port))
    t.connect(username=username, password=password)
    sftp = paramiko.SFTPClient.from_transport(t)
    sftp.get(server_path, local_path)
    t.close()

def _mkdirs(sftp, path: str, sep="/"):
    '''
    级联创建指定目录，
    :param sftp:sftp实例对象
    :param path:目录路径
    :param sep: 远程服务器路径分隔符
    :return:
    '''
    re_split = path.split(sep)
    re_multi = ["/".join(re_split[:i + 1]) for i in range(len(re_split))]  # 生成逐层路径列表
    for i in re_multi:
        if not i:
            continue
        f_list = sftp.listdir_attr(os.path.dirname(i))
        # 起始字符为d均为目录
        dir_list = [str(i).split()[-1] for i in f_list if str(i).startswith("d")]
        if os.path.basename(i) not in dir_list:
            sftp.mkdir(i)


def sftp_upload_dir(host_ip, username, password, remote_dir, local_dir, update_flag=False, port=22, key_file=None, **kwargs):
    '''
    向远程linux上传文件夹
    :param host_ip: 主机IP
    :param username: 用户名
    :param password: 密码
    :param remote_dir: 远程linux文件夹路径，支持创建不存在的路径
    :param local_dir: 待上传的本地文件夹路径
    :param update_flag: 删除标志，True--做覆盖操作，False--文件已存在时保留
    :param port: ftp端口
    :param kwargs:
    :return:
    '''
    # private_key = paramiko.RSAKey.from_private_key_file(key_file) # 实例化一个私钥对象
    # transport = paramiko.Transport((host_ip, port))
    # transport.connect(username, pkey=private_key)
    # sftp = paramiko.SFTPClient.from_transport(transport)
    #TODO
    #没调通
    sep = "/"
    t = paramiko.Transport((host_ip, port))
    t.connect(username=username, password=password)
    sftp = paramiko.SFTPClient.from_transport(t)

    _mkdirs(sftp, remote_dir)  # 级联创建远程目录
    remote_path = remote_dir.rstrip(sep)
    logger.info(f"{local_dir}, {remote_path}")
    for root, dirs, files in os.walk(local_dir):
        # 遍历上传目录下文件
        for f in files:
            local_file_path = os.path.join(root, f)
            remote_file_dir = remote_path + sep + root.replace(local_dir, '').replace(os.sep,sep)
            _sep = sep if not remote_file_dir.endswith(sep) else ""
            remote_file_path = remote_file_dir + _sep + f
            logger.info(f'local_file_path: {local_file_path}  remote_file_path:{remote_file_path}')
            sftp.put(local_file_path, remote_file_path)
        # 遍历创建子目录
        for d in dirs:
            # full_dir = os.path.join(root, d)
            full_dir = os.path.join(root, d)
            sub_dir = full_dir.replace(local_dir,'').replace(os.sep, sep)
            remote_dir = remote_path + sub_dir
            logger.info(remote_dir)
            # 排除父目录中已存在的子目录防止目录创建冲突
            f_list = sftp.listdir_attr(os.path.dirname(remote_dir))
            dir_list = [str(i).split()[-1] for i in f_list if str(i).startswith("d")]
            if os.path.basename(remote_dir) not in dir_list:
                sftp.mkdir(remote_dir)
    t.close()


def sftp_upload_dir_to_ac(host_ip, username, remote_dir, local_dir, port=22, key_file=None, **kwargs):
    """
    把本地代码上传至ac  --创建目录会有权限问题，待解决
    :param host_ip:
    :param username:
    :param remote_dir:
    :param local_dir:
    :param port:
    :param key_file:
    :param kwargs:
    :return:
    """
    sep = "/"  # linux目录分隔符
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=host_ip, port=port, username=username, pkey=key_file)
    sftp = ssh.open_sftp()

    # AC创建目录提示权限拒绝
    # _mkdirs(sftp, remote_dir)  # 级联创建远程目录
    remote_path = remote_dir.rstrip(sep)
    logger.info(f"{local_dir}, {remote_path}")
    for root, dirs, files in os.walk(local_dir):
        # 遍历上传目录下文件
        for f in files:
            if f.endswith('.pyc') or f.startswith('.'):  # 过滤掉编译文件
                continue
            local_file_path = os.path.join(root, f)
            remote_file_dir = remote_path + sep + root.replace(local_dir, '').replace(os.sep,sep)
            _sep = sep if not remote_file_dir.endswith(sep) else ""
            remote_file_path = remote_file_dir + _sep + f
            logger.info(f'local_file_path: {local_file_path}\nremote_file_path:{remote_file_path}')
            sftp.put(local_file_path, remote_file_path)
        # 遍历创建子目录
        for d in dirs:
            full_dir = os.path.join(root, d)
            for _fn in full_dir.split(os.path.sep): # 过滤.idea等目录
                if _fn.startswith('.'):
                    continue
            else:
                sub_dir = full_dir.replace(local_dir,'').replace(os.sep, sep)
                remote_dir = remote_path + sub_dir
                logger.info(remote_dir)
                # 排除父目录中已存在的子目录防止目录创建冲突
                f_list = sftp.listdir_attr(os.path.dirname(remote_dir))
                dir_list = [str(i).split()[-1] for i in f_list if str(i).startswith("d")]
                if os.path.basename(remote_dir) not in dir_list:
                    sftp.mkdir(remote_dir)
    sftp.close()
    ssh.close()



def ssh_command_after_connect_to_target_via_jump_host(jump_host, jump_port, jump_user, jump_password, target_host, target_port,
                                                      target_user, target_password, command="pwd;ls -l"):
    """
    通过跳板机登陆目标节点
    :param jump_host:
    :param jump_port:
    :param jump_user:
    :param jump_password:
    :param target_host:
    :param target_port:
    :param target_user:
    :param target_password:
    :param command:
    :return:
    """
    # 创建 SSH 客户端
    jump_client = paramiko.SSHClient()
    jump_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    # 连接到跳板机
    jump_client.connect(jump_host, port=jump_port, username=jump_user, password=jump_password)
    print("Connected to jump host")
    stdin, stdout, stderr = jump_client.exec_command(command)
    print(stdout.read().decode())

    # 创建一个通道
    jump_transport = jump_client.get_transport()
    dest_addr = (target_host, target_port)
    # local_addr = (jump_host, jump_port)
    local_addr = ('localhost', 0)


    # 建立到目标主机的连接
    channel = jump_transport.open_channel("direct-tcpip", dest_addr, local_addr)

    # 创建目标主机的 SSH 客户端
    target_client = paramiko.SSHClient()
    target_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    # 使用通道连接目标主机
    target_client.connect(target_host, port=target_port, username=target_user, password=target_password,
                          sock=channel)
    print("Connected to target host")

    # 在目标主机上执行命令
    stdin, stdout, stderr = target_client.exec_command(command)
    print(stdout.read().decode())

    # 关闭连接
    target_client.close()
    jump_client.close()

if __name__ == '__main__':
    # AC昆山登陆节点信息
    user=GLO_VAR["AC_KS"]["username"]
    ip=GLO_VAR["AC_KS"]["host_ip"]
    port=GLO_VAR["AC_KS"]["port"]
    key_path = GLO_VAR["AC_KS"]["IdentityFile"]
    local_dir = r"C:\Users\liuxu\Desktop\py-auto-test\profiler_func_auto_testproj"
    remote_dir = r"/public/home/sugonhpctest01/home/liuxu/auto_test/profiler_func_auto_testproj/"


    # os.system(r"C:\Users\liuxu\Desktop\py-auto-test\profiler_func_auto_testproj\clean_cache.bat")
    # sftp_upload_dir_to_ac(host_ip=ip, username=user, password=None, remote_dir=remote_dir, local_dir=local_dir, port=port, key_file=key_path)
    sftp_upload_dir_to_ac(host_ip=ip, username=user, password=None, remote_dir=remote_dir, local_dir=local_dir,
                          port=port,
                          key_file=key_path)

    # 上传目录至乌镇裸金属
    # sftp_upload_dir(host_ip="10.2.4.48", username="root", password="Newsys0!", remote_dir="/construct02/home/liuxu/prof_auto_test", local_dir=local_dir, port=22)

    # 上传到本地虚拟机 --测试ok
    # sftp_upload_dir(host_ip="192.168.22.139", username="root", password="root", remote_dir="/home/auto_test", local_dir=local_dir, port=22)



    # # ssh AC登陆节点  --测试ok
    # ssh_command(host_ip=ip, username=user, password=None, port=port, key_filename=key_path, command="pwd;ls")

    # # 登陆AC计算节点 --会抛出异常Secsh channel 1 open FAILED: Connection refused: Connect failed
    # connect_to_target_via_jump_host(jump_host=ip, jump_port=port, jump_user=user, jump_password=None, target_host="b13r2n08", target_port=None, target_user=user, target_password=None )




    # ssh_command(host_ip="10.2.4.43", username="root", password="SUgOn$WuzHeN!miX_2022%6+17\bJD", command="ls")
    # commands = ["rpm -qa | grep rock","df -h", "rpm -V rock-5.7.1"]
    # ssh_commands(host_ip="10.2.4.48", username="root", password="Newsys0!", commands=[cmd])















